import os
import re
import sys
import glob
import shutil
import random

parent = os.path.dirname(os.path.dirname(__file__))
sys.path.append(parent)
from util_my import github_urls

im_ext = ['jpg', 'png', 'bmp']

def setup(git_name='', github_dir=''):
    
    repo_url = github_urls[git_name]
    git_dir = os.path.join(github_dir, repo_url.split('/')[-1])

    if os.path.exists(git_dir):
        # TODO remove git to code
        downloads_dir = os.path.join(parent, 'downloads')
        code_dir = os.path.join(downloads_dir, git_name)
        os.makedirs(downloads_dir, exist_ok=True)

        if os.path.exists(code_dir):
            shutil.rmtree(code_dir)

        for src in glob.glob(os.path.join(git_dir, '**', '**.py'), recursive=True) \
                + glob.glob(os.path.join(git_dir, '**', '**.jpg'), recursive=True) \
                + glob.glob(os.path.join(git_dir, '**', '**.yaml'), recursive=True):

            dst = src.replace(git_dir, code_dir)
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            shutil.copy(src, dst)

        return git_dir, code_dir
    else:
        return git_dir, ''

def move_files(split_dir, file:str, data_dir, exts=im_ext+['xml', 'txt', 'json']):
    os.makedirs(split_dir, exist_ok=True)
    try:
        for ext in exts:
            path = split_dir + re.sub('\w+$', ext, file).replace(data_dir, '')
            if not os.path.exists(path) and os.path.exists(re.sub('\w+$', ext, file)):
                os.makedirs(os.path.dirname(path), exist_ok=True)
                shutil.move(re.sub('\w+$', ext, file), path)
    except:
        pass

def get_images(dir_):
    images = []
    for ext in im_ext:
        images += glob.glob(os.path.join(dir_, '**', f'**.{ext}'), recursive=True)
    return images

def split_change(len_im, split, len_split):
    return len_im * split == 0 or (int(len_im * split) >= len_split - 1 and int(len_im * split) <= len_split + 1)

def update_labels(data_dir='', split='1::'):

    total = sum([int(s) for s in split.split(':') if s])

    train_count = split.split(':')[0]
    val_count = split.split(':')[1]
    test_count = split.split(':')[2]

    train_dir = os.path.join(data_dir, 'train')
    val_dir = os.path.join(data_dir, 'val' if val_count else 'train')
    test_dir = os.path.join(data_dir, 'test' if test_count else 'train')

    train_split = int(train_count) / total
    val_split = (int(val_count) if val_count else 0) / total
    test_split = (int(test_count) if test_count else 0) / total

    train_images = get_images(train_dir)
    val_images = get_images(val_dir)
    test_images = get_images(test_dir)

    images = get_images(data_dir)

    if split_change(len(images), train_split, len(train_images)) and \
        split_change(len(images), test_split, len(test_images)) and \
        split_change(len(images), val_split, len(val_images)):
        return train_dir, val_dir

    random.shuffle(images)

    for i, image in enumerate(images):
        if i <= len(images) * train_split:
            move_files(train_dir, image, data_dir)
        elif i <= len(images) * (train_split + val_split):
            move_files(val_dir, image, data_dir)
        elif i <= len(images) * (train_split + val_split + test_split):
            move_files(test_dir, image, data_dir)
    
    for fd in os.listdir(data_dir):
        if fd not in ['train', 'test', 'val']:
            shutil.rmtree(os.path.join(data_dir, fd))
    
    return train_dir, val_dir

def set_coco_image(labels, id, width, height, file_name):
    '''
    ```python
    # https://cocodataset.org/#format-data
    image = dict(
        file_name=os.path.basename('1.jpg'),
        height=640,
        width=480,
        id=0
    )
    ```
    '''
    image = dict(
        id=id,
        width=width,
        height=height,
        file_name=file_name,
    )
    labels['images'].append(image)

    return id

def set_coco_category(labels, name, num_keypoints):
    '''
    '''

    if len([n for n in labels['categories'] if n['name'] == name]) == 0:
        category = dict(
            name=name,
            id=len(labels['categories']),
            keypoints=[f'keypoints{i}' for i in list(range(num_keypoints))],
            skeleton=[[i, i+1] for i in list(range(num_keypoints-1))] + [[num_keypoints-1, 0]],
        )
        labels['categories'].append(category)

    category_id = labels['categories'].index([n for n in labels['categories'] if n['name'] == name][0])

    return category_id

def set_coco_annotation(labels, id, width, height, keypoints, image_id, name):
    '''
    x1 = y1 = x2 = y2 = x3 = y3 = 0
    seg = [x1, y1, x2, y2, x3, y3]

    annotation = dict(
        segmentation=[seg, seg],
        num_keypoints=len(),
    )

    labels = dict(
        images=[image],
        annotations=[annotation],
    )
    '''

    num_keypoints = len(keypoints) // 3
    category_id = set_coco_category(labels, name, num_keypoints)
    
    annotation = dict(
        bbox=[],
        iscrowd=0,
        area=width*height,
        image_id=image_id,
        keypoints=keypoints,
        num_keypoints=num_keypoints,
        id=int(str(image_id)+str(id)),
        category_id=category_id,
    )

    labels['annotations'].append(annotation)

def labelme2coco(label_dir='', out_dir='', categories=[]):

    import json
    import tqdm

    labels = dict(
        images=[],
        annotations=[],
        categories=categories,
    )

    for id, json_file in enumerate(tqdm.tqdm(glob.glob(os.path.join(label_dir, '*.json')))):
        label = json.load(open(json_file, encoding='utf-8'))

        try:
            width = label['imageWidth']
            height = label['imageHeight']
            file_name = label['imagePath'].split(os.sep)[-1].split('/')[-1]

            if not os.path.exists(os.path.join(label_dir, file_name)):
                continue
    
            image_id = set_coco_image(labels, id, width, height, file_name)

            for id, shape in enumerate(label['shapes']):                
                name = shape['label']
                points = shape['points']
                keypoints = []
                for point in points:
                    keypoints.append(int(point[0]))
                    keypoints.append(int(point[1]))
                    keypoints.append(2)

                set_coco_annotation(labels, id, width, height, keypoints, image_id, name)

        except Exception as e:
            print(e)

    json.dump(labels, open(os.path.join(out_dir, 'data.json'), 'w'))
    
    return labels

def labelimg2coco(label_dir='', out_dir='', categories=[]):
    import json
    import tqdm

    labels = dict(
        images=[],
        annotations=[],
        categories=categories,
    )

    from xml.dom.minidom import Element, Node, CDATASection, parse

    for id, xml_path in enumerate(tqdm.tqdm(glob.glob(os.path.join(label_dir, '*.xml')))):

        tree:Element = parse(xml_path)

        def get_data(tree:Element, name):
            node:Node = tree.getElementsByTagName(name)[0]
            section:CDATASection = node.childNodes[0]
            return section.data
        try:
            width = int(get_data(tree, 'width'))
            height = int(get_data(tree, 'height'))
            file_name = get_data(tree, 'filename')
        except:
            print('Could not parse size in ', xml_path)
            return
        
        if not os.path.exists(os.path.join(label_dir, file_name)):
            continue

        image_id = set_coco_image(labels, id, width, height, file_name)

        for obj in tree.getElementsByTagName('objects'):
            name = get_data(obj, 'name')
            points = str(get_data(obj, 'points')).split(';')
            keypoints = []
            for point in points:
                keypoints.append(int(float(point.split(',')[0])))
                keypoints.append(int(float(point.split(',')[1])))
                keypoints.append(2)

            set_coco_annotation(labels, id, width, height, keypoints, image_id, name)

    json.dump(labels, open(os.path.join(out_dir, 'data.json'), 'w'))
    
    return labels

if __name__ == '__main__':
    labelimg2coco(r'C:\Users\sv00296941\Downloads\20230302112923ai', r'D:\Docker\daai\target\data')